package algorithm.graph;

import java.awt.image.SampleModel;
import java.util.*;

public class Kruskal {

    public static class UnionFind{
        public HashMap<Node, Node> parents;
        public HashMap<Node, Integer> sizeMap;


        public void makeSets(Collection<Node> nodes){
            parents.clear();
            sizeMap.clear();
            for (Node node : nodes) {
                parents.put(node, node);
                sizeMap.put(node, 1);
            }
        }

        public Node findFather(Node node){
            Stack<Node> stack = new Stack<>();
            while(parents.get(node) != node){
                stack.push(node);
                node = parents.get(node);
            }

            while(!stack.isEmpty()){
                Node pop = stack.pop();
                parents.put(pop, node);
            }

            return node;
        }


        public boolean isSameSet(Node node1, Node node2){
            return findFather(node1) == findFather(node2);
        }

        public void union(Node a, Node b){
            Node aFather = findFather(a);
            Node bFather = findFather(b);
            if(aFather == bFather)
                return;
            Integer aSize = sizeMap.get(aFather);
            Integer bSize = sizeMap.get(bFather);
            Node big = aSize > bSize? aFather : bFather;
            Node small = big == aFather? bFather : aFather;
            parents.put(small, big);
            sizeMap.put(big, aSize + bSize);
            sizeMap.remove(small);
        }
    }

    public static Set<Edge> kruskalMST(Graph graph){

        HashSet<Edge> set = new HashSet<>();
        UnionFind unionFind = new UnionFind();
        unionFind.makeSets(graph.nodes.values());
        PriorityQueue<Edge> queue = new PriorityQueue<>(new Comparator<Edge>() {
            @Override
            public int compare(Edge o1, Edge o2) {
                return o1.weight - o2.weight;
            }
        });
        for (Edge edge : graph.edges) {
            queue.add(edge);
        }
        while(!queue.isEmpty()){
            Edge edge = queue.poll();
            if(!unionFind.isSameSet(edge.from, edge.to)){
                set.add(edge);
                unionFind.union(edge.from, edge.to);
            }
        }
        return set;
    }
}
